//
// Created by DonAzufre on 24-5-19.
//

#include "llvm/Pass.h"
#include "llvm/IR/Function.h"
#include "llvm/Support/raw_ostream.h"
#include <set>

using namespace llvm;

namespace {
    using namespace std;

    struct CFGEdgeNum : public FunctionPass {
        static char ID;
        set<const BasicBlock*> visited;

        CFGEdgeNum() : FunctionPass(ID) {}

        void dfsBB(const BasicBlock* bb, int& edgeNum) {
            const Instruction* last = nullptr;
            for (auto &i : *bb) {
                if (i.isTerminator()) {
                    last = &i;
                    break;
                }
            }

            assert(last);

            if (last->getNumSuccessors() == 0)
                return;
            else {
                edgeNum += last->getNumSuccessors();
                visited.insert(bb);
                for (int i = 0; i < last->getNumSuccessors(); i++) {
                    if (auto succ = last->getSuccessor(i); visited.find(succ) == visited.end()) {
                        dfsBB(succ, edgeNum);
                    }
                }
            }

        }

        bool runOnFunction(Function &F) override {

            int edges = 0;
            visited.clear();
            dfsBB(&F.getEntryBlock(), edges);
            errs() << "[CFGEdgeNum] In function " << F.getName() << ", number of edges in cfg is " << edges << '\n';

            return false;
        }
    };
}

char CFGEdgeNum::ID = 0;
static RegisterPass<CFGEdgeNum> X("cfg-edge-num", "Count edges of cfg of each function");